import socket
import cv2
import av
import struct
import time
from io import BytesIO

BUFFER_SIZE = 4194304  # 1MB缓冲区
SERVER_IP = '0.0.0.0'
SERVER_PORT = 6001

def draw_hud(img, fps, bandwidth, fps_history, bandwidth_history):
    """在图像上绘制HUD信息"""
    # 在右上角显示文本信息
    text = f"FPS: {fps:.1f}  BW: {bandwidth:.2f} MB/s"
    text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
    text_x = img.shape[1] - text_size[0] - 10
    text_y = text_size[1] + 10
    cv2.putText(img, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

    # 绘制统计图表区域
    chart_x = img.shape[1] - 210
    chart_y = text_y + 20
    chart_width = 200
    chart_height = 80

    # 绘制背景
    cv2.rectangle(img, (chart_x, chart_y), 
                 (chart_x + chart_width, chart_y + chart_height),
                 (0, 0, 0), -1)

    # 绘制帧率折线（绿色）
    if len(fps_history) > 1:
        max_fps = max(fps_history) or 1
        for i in range(1, len(fps_history)):
            x1 = chart_x + int((i-1) * (chart_width / (len(fps_history)-1)))
            y1 = chart_y + int(chart_height - (fps_history[i-1]/max_fps * chart_height))
            x2 = chart_x + int(i * (chart_width / (len(fps_history)-1)))
            y2 = chart_y + int(chart_height - (fps_history[i]/max_fps * chart_height))
            cv2.line(img, (x1, y1), (x2, y2), (0, 255, 0), 1)

    # 绘制带宽折线（蓝色）
    if len(bandwidth_history) > 1:
        max_bw = max(bandwidth_history) or 1
        for i in range(1, len(bandwidth_history)):
            x1 = chart_x + int((i-1) * (chart_width / (len(bandwidth_history)-1)))
            y1 = chart_y + int(chart_height - (bandwidth_history[i-1]/max_bw * chart_height))
            x2 = chart_x + int(i * (chart_width / (len(bandwidth_history)-1)))
            y2 = chart_y + int(chart_height - (bandwidth_history[i]/max_bw * chart_height))
            cv2.line(img, (x1, y1), (x2, y2), (255, 0, 0), 1)

    return img

def main():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((SERVER_IP, SERVER_PORT))
    sock.listen(1)
    
    codec = av.CodecContext.create('hevc', 'r')
    
    while True:
        conn, addr = sock.accept()
        print(f"连接来自: {addr}")
        
        # 接收配置包
        config_data = b''
        while len(config_data) < 12:
            chunk = conn.recv(12 - len(config_data))
            if not chunk:
                break
            config_data += chunk
        if len(config_data) != 12 or struct.unpack("<L", config_data[:4])[0] != 0xFFFFFFFF:
            print("配置包错误")
            conn.close()
            continue
        width, height = struct.unpack("<LL", config_data[4:12])
        print(f"新分辨率: {width}x{height}")
        codec = av.CodecContext.create('hevc', 'r')
        
        # 初始化统计变量
        start_time = time.time()
        frame_counter = 0
        bytes_received_total = 0
        current_fps = 0
        current_bandwidth = 0.0
        fps_history = []
        bandwidth_history = []
        
        try:
            while True:
                # 接收帧头
                header = b''
                while len(header) < 8:
                    chunk = conn.recv(8 - len(header))
                    if chunk:
                        bytes_received_total += len(chunk)
                    if not chunk:
                        break
                    header += chunk
                if len(header) < 8:
                    break
                total_size, nalu_count = struct.unpack("<LL", header)
                
                # 接收帧数据
                received = 0
                buffer = BytesIO()
                while received < total_size:
                    chunk = conn.recv(min(BUFFER_SIZE, total_size - received))
                    if chunk:
                        bytes_received_total += len(chunk)
                    if not chunk:
                        break
                    buffer.write(chunk)
                    received += len(chunk)

                # 更新统计信息
                current_time = time.time()
                if current_time - start_time >= 1.0:
                    current_fps = frame_counter
                    current_bandwidth = bytes_received_total / 1e6  # 转换为MB/s
                    fps_history.append(current_fps)
                    bandwidth_history.append(current_bandwidth)
                    # 保持最近30秒数据
                    if len(fps_history) > 30:
                        fps_history.pop(0)
                    if len(bandwidth_history) > 30:
                        bandwidth_history.pop(0)
                    frame_counter = 0
                    bytes_received_total = 0
                    start_time = current_time

                # 分割NAL单元并解码
                buffer.seek(0)
                raw_data = buffer.getvalue()
                pos = 0
                
                for _ in range(nalu_count):
                    nalu_size = struct.unpack("<L", raw_data[pos:pos+4])[0]
                    pos += 4
                    nalu_data = raw_data[pos:pos+nalu_size]
                    pos += nalu_size
                    
                    packet = av.packet.Packet(nalu_data)
                    try:
                        frames = codec.decode(packet)
                        for frame in frames:
                            if isinstance(frame, av.VideoFrame):
                                img = frame.to_ndarray(format='bgr24')
                                frame_counter += 1
                                
                                # 保存原始图像用于截图
                                current_frame = img.copy()
                                
                                # 绘制HUD
                                img = draw_hud(img, current_fps, current_bandwidth,
                                             fps_history, bandwidth_history)
                                
                                # 显示图像
                                cv2.imshow('H265 Stream', img)
                                
                                # 处理键盘事件
                                key = cv2.waitKey(1)
                                if key == 27:  # ESC退出
                                    raise KeyboardInterrupt
                                elif key == ord('s'):  # 截图
                                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                                    filename = f"screenshot_{timestamp}.png"
                                    cv2.imwrite(filename, current_frame)
                                    print(f"截图已保存: {filename}")
                                
                    except Exception as e:
                        print(f"解码错误: {e}")
                
                buffer = BytesIO()
                
        except (ConnectionResetError, KeyboardInterrupt):
            print("连接断开")
        finally:
            conn.close()
            cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
